# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Common functions
"""

from typing import Union, List, Tuple, Iterable
from datetime import time, timedelta, date
import numpy as np
from numpy import ndarray
import torch
from torch import Tensor
from torch.nn import Parameter


__all__ = [
    'PI',
    'keepdims_sum',
    'keepdims_mean',
    'keepdims_prod',
    'reduce_any',
    'reduce_all',
    'reduce_prod',
    'concat_first_dim',
    'concat_last_dim',
    'concat_penulti',
    'stack_first_dim',
    'stack_last_dim',
    'stack_penulti',
    'squeeze_first_dim',
    'squeeze_last_dim',
    'squeeze_penulti',
    'identity',
    'periodic_variable',
    'periodic_difference',
    'gather_vector',
    'gather_value',
    'pbc_box_reshape',
    'pbc_image',
    'coordinate_in_pbc',
    'vector_in_pbc',
    'calc_vector_nopbc',
    'calc_vector_pbc',
    'calc_vector',
    'calc_distance_nopbc',
    'calc_distance_pbc',
    'calc_distance',
    'calc_angle_by_vectors',
    'calc_angle_nopbc',
    'calc_angle_pbc',
    'calc_angle',
    'calc_torsion_by_vectors',
    'calc_torsion_nopbc',
    'calc_torsion_pbc',
    'calc_torsion',
    'coulomb_interaction',
    'lennard_jones_potential',
    'lennard_jones_potential2',
    'get_integer',
    'get_ndarray',
    'get_tensor',
    'get_ms_array',
    'check_broadcast',
    'any_none',
    'all_none',
    'any_not_none',
    'all_not_none',
    'get_arguments',
    'get_initializer',
    'bonds_in',
    'GLOBAL_DEVICE_SPONGE',
    'set_torchsponge_global_device',
    'toxyz'
]

PI = 3.141592653589793238462643383279502884197169399375105820974944592307
r""":math:`\pi`"""


def keepdims_sum(x: Tensor, axis: Union[int, Tuple[int], list[int]] = ()) -> Tensor:

    return torch.sum(x, dim=axis, keepdim=True)


def keepdims_mean(x: Tensor, axis: Union[int, Tuple[int], list[int]] = ()) -> Tensor:

    return torch.mean(x, dim=axis, keepdim=True)


def keepdims_prod(x: Tensor, axis: Union[int, Tuple[int], list[int]] = ()) -> Tensor:

    return torch.prod(x, dim=axis, keepdim=True)


def reduce_any(x: Tensor, axis: Union[int, Tuple[int], list[int]] = ()) -> Tensor:

    return torch.any(x, dim=axis)


def reduce_all(x: Tensor, axis: Union[int, Tuple[int], list[int]] = ()) -> Tensor:

    return torch.all(x, dim=axis)


def reduce_prod(x: Tensor, axis: Union[int, Tuple[int], List[int]] = ()) -> Tensor:

    return torch.prod(x, dim=axis, keepdim=False)


def concat_first_dim(input_x: Union[tuple, list]) -> Tensor:
   
    return torch.cat(input_x, dim=0)


def concat_last_dim(input_x: Union[tuple, list]) -> Tensor:
    
    return torch.cat(input_x, dim=-1)


def concat_penulti(input_x: Union[tuple, list]) -> Tensor:

    return torch.cat(input_x, dim=-2)


def stack_first_dim(input_x: Union[tuple, list]) -> Tensor:

    return torch.stack(input_x, dim=0)


def stack_last_dim(input_x: Union[tuple, list]) -> Tensor:
  
    return torch.stack(input_x, dim=-1)


def stack_penulti(input_x: Union[tuple, list]) -> Tensor:
    
    return torch.stack(input_x, dim=-2)


def squeeze_first_dim(input_x: Tensor) -> Tensor:

    return torch.squeeze(input_x, 0)


def squeeze_last_dim(input_x: Tensor) -> Tensor:
    
    return torch.squeeze(input_x, -1)


def squeeze_penulti(input_x: Tensor) -> Tensor:
    
    return torch.squeeze(input_x, -2)


def identity(x: Tensor) -> Tensor:

    return torch.nn.Identity()(x)


def periodic_variable(variable: Tensor,
                      upper: Tensor,
                      lower: Tensor = 0,
                      mask: Tensor = None,
                      ) -> Tensor:

    period = upper - lower
    period_value = variable - period * torch.floor((variable - lower) / period)
    if mask is None:
        return period_value
    if mask.shape != variable.shape:
        mask = torch.broadcast_to(mask, variable.shape)
    return torch.where(mask, period_value, variable)


def periodic_difference(difference: Tensor,
                        period: Tensor,
                        mask: Tensor = None,
                        offset: float = -0.5,
                        ) -> Tensor:
    r"""get value of difference between periodic variables.

    """

    period_diff = difference - period * torch.floor(difference / period - offset)
    if mask is None:
        return period_diff
    if mask.shape != difference.shape:
        mask = torch.broadcast_to(mask, difference.shape)
    return torch.where(mask, period_diff, difference)


def gather_vector(tensor: Tensor, index: Tensor) -> Tensor:
    r"""Gather vector from the penultimate axis (`axis=-2`) of the tensor according to index.

    """
    # index = index.to(torch.int64)
    if index.shape[0] == 1:
        return torch.index_select(tensor, -2, index[0].view(-1)).reshape(tensor.shape[:-2]+index[0].shape[0:]+tensor.shape[-1:])
    if tensor.shape[0] == 1:
        return torch.index_select(tensor[0], -2, index.view(-1)).reshape(tensor[0].shape[:-2]+index.shape[0:]+tensor[0].shape[-1:])

    # # (B, N, M)
    # shape0 = index.shape
    # # (B, N * M, 1) <- (B, N, M)
    # index = torch.reshape(index, (shape0[0], -1, 1))
    # # (B, N * M, D) <- (B, N, D)
    # vectors = torch.gather(tensor, -2, index[0])
    # # (B, N, M, D) <- (B, N, M) + (D,)
    # output_shape = shape0 + tensor.shape[-1:]
    # # (B, N, M, D)
    # return torch.reshape(vectors, output_shape)
    shape = index.shape+(tensor.shape[-1],)
    index = index.unsqueeze(-1).expand(*shape)
    tensor = tensor.unsqueeze(-2).expand(*shape)
    output = torch.gather(tensor, 1, index)

    return output

def gather_value(tensor: Tensor, index: Tensor) -> Tensor:
    r"""Gather value from the last axis (`axis=-1`) of the tensor according to index.

    """

    if index.shape[0] == 1:
        return torch.index_select(tensor, -1, index[0].view(-1)).reshape(tensor.shape[:-1]+index[0].shape[0:])
    if tensor.shape[0] == 1:
        return torch.index_select(tensor[0], -1, index.view(-1)).reshape(tensor[0].shape[:-1]+index.shape[0:])

    # # (B, N, M)
    # origin_shape = index.shape
    # # (B, N * M) <- (B, N, M)
    # index = torch.reshape(index, (origin_shape[0], -1))
    # # (B, N * M)
    # values = torch.gather(tensor, -1, index)
    # (B, N, M)
    # return torch.reshape(values, origin_shape)
    # output = torch.index_select(tensor, -1, index.view(-1)).reshape(tensor.shape[:-1]+index.shape[0:])
    shape = tensor.shape + (index.shape[-1],)
    output = torch.gather(tensor.unsqueeze(-1).expand(*shape), 1, index)
    return output


def pbc_box_reshape(pbc_box: Tensor, ndim: int) -> Tensor:
    r"""Reshape the pbc_box as the same ndim.

    """

    if ndim <= 2:
        return pbc_box
    shape = pbc_box.shape[:1] + (1,) * (ndim - 2) + pbc_box.shape[-1:]
    return torch.reshape(pbc_box, shape)

def pbc_image(position: Tensor, pbc_box: Tensor, offset: float = 0) -> Tensor:
    r"""calculate the periodic image of the PBC box

    """

    pbc_box = pbc_box_reshape(pbc_box.detach(), position.ndim)
    image = -torch.floor(position / pbc_box - offset)
    return image.to(torch.int32)

def coordinate_in_pbc(position: Tensor, pbc_box: Tensor, offset: float = 0) -> Tensor:
    r"""get coordinate in main PBC box
    """

    pbc_box = pbc_box_reshape(pbc_box.detach(), position.ndim)
    return position - pbc_box * torch.floor(position / pbc_box - offset)

def vector_in_pbc(vector: Tensor, pbc_box: Tensor, offset: float = -0.5) -> Tensor:
    r"""Make the value of vector :math:`\vec{v}` at a single PBC box :math:`\vec{L}`.

    """

    pbc_box = pbc_box_reshape(pbc_box, vector.ndim)
    box_nograd = pbc_box.detach()
    inv_box = torch.reciprocal(box_nograd)
    vector = vector - box_nograd * torch.floor(vector * inv_box - offset)
    return vector * inv_box * pbc_box

def calc_vector_nopbc(initial: Tensor, terminal: Tensor) -> Tensor:
    r"""Compute vector from initial point to terminal point without perodic bundary condition.

    """

    return terminal - initial

def calc_vector_pbc(initial: Tensor, terminal: Tensor, pbc_box: Tensor) -> Tensor:
    r"""Compute vector from initial point to terminal point at perodic bundary condition.

    """

    return vector_in_pbc(terminal-initial, pbc_box)

def calc_vector(initial: Tensor, terminal: Tensor, pbc_box: Tensor = None) -> Tensor:
    r"""Compute vector from initial point to terminal point.

    """

    vector = terminal - initial
    if pbc_box is None:
        return vector
    return vector_in_pbc(vector, pbc_box)

def calc_distance_nopbc(position_a: Tensor,
                        position_b: Tensor,
                        keepdim: bool = False,
                        ) -> Tensor:
    r"""Compute distance between position A and B without perodic bundary condition.


    """

    vec = calc_vector_nopbc(position_a, position_b)
    return torch.norm(vec, dim=-1, keepdim=keepdim)

def calc_distance_pbc(position_a: Tensor,
                      position_b: Tensor,
                      pbc_box: Tensor = None,
                      keepdim: bool = False
                      ) -> Tensor:
    r"""Compute distance between position :math:`A` and :math:`B` at perodic bundary condition.

    """

    vec = calc_vector_pbc(position_a, position_b, pbc_box)
    return torch.norm(vec, dim=-1, keepdim=keepdim)

def calc_distance(position_a: Tensor,
                  position_b: Tensor,
                  pbc_box: Tensor = None,
                  keepdim: bool = False,
                  ) -> Tensor:
    r"""Compute distance between position :math:`A` and :math:`B`.

    """

    vec = calc_vector_nopbc(position_a, position_b)
    if pbc_box is not None:
        vec = vector_in_pbc(vec, pbc_box)
    return torch.norm(vec, dim=-1, keepdim=keepdim)

def calc_angle_by_vectors(vector1: Tensor,
                          vector2: Tensor,
                          keepdim: bool = False
                          ) -> Tensor:
    r"""Compute angle between two vectors.

    """

    # (...) or (..., 1) <- (..., D)
    dis1 = torch.norm(vector1, axis=-1, keepdim=keepdim)
    dis2 = torch.norm(vector2, axis=-1, keepdim=keepdim)
    dot12 = torch.sum(vector1*vector2, axis=-1, keepdim=keepdim)
    # (...) or (..., 1)
    cos_theta = dot12 / dis1 / dis2
    return torch.acos(cos_theta)

def calc_angle_nopbc(position_a: Tensor,
                     position_b: Tensor,
                     position_c: Tensor,
                     keepdim: bool = False,
                     ) -> Tensor:
    r"""Compute angle :math:`\angle{ABC}` formed by the position coordinates of three positions
    
    """

    # (...,D)
    vec_ba = calc_vector_nopbc(position_b, position_a)
    vec_bc = calc_vector_nopbc(position_b, position_c)
    return calc_angle_by_vectors(vec_ba, vec_bc, keepdim=keepdim)

def calc_angle_pbc(position_a: Tensor,
                   position_b: Tensor,
                   position_c: Tensor,
                   pbc_box: Tensor,
                   keepdim: bool = False,
                   ) -> Tensor:
    r"""Compute angle :math:`\angle{ABC}` formed by the position coordinates of three positions

    """

    # (B, ..., D)
    vec_ba = calc_vector_pbc(position_b, position_a, pbc_box)
    vec_bc = calc_vector_pbc(position_b, position_c, pbc_box)
    return calc_angle_by_vectors(vec_ba, vec_bc, keepdim=keepdim)

def calc_angle(position_a: Tensor,
               position_b: Tensor,
               position_c: Tensor,
               pbc_box: Tensor = None,
               keepdim: bool = False,
               ) -> Tensor:
    r"""Compute angle formed by three positions :math:`A`, :math:`B` and :math:`C`
        with or without periodic boundary condition.


    """

    if pbc_box is None:
        return calc_angle_nopbc(position_a, position_b, position_c, keepdim=keepdim)
    return calc_angle_pbc(position_a, position_b, position_c, pbc_box=pbc_box, keepdim=keepdim)

def calc_torsion_by_vectors(vector1: Tensor,
                            vector2: Tensor,
                            axis_vector: Tensor = None,
                            keepdim: bool = False,
                            ) -> Tensor:
    r"""Compute torsion angle formed by two direction vectors :math:`\vec{v_1}` and :math:`\vec{v_2}`
        and an axis vector :math:`\vec{v_{axis}}`.

    """

    if axis_vector is None:
        return calc_angle_by_vectors(vector1, vector2, keepdim=keepdim)

    # (..., D)
    vec_a = torch.cross(vector1, axis_vector)
    vec_b = torch.cross(vector2, axis_vector)
    cross_ab = torch.cross(vec_a, vec_b)

    # (..., 1) <- (..., D)
    axis_norm = torch.norm(axis_vector, dim=-1, keepdim=True)
    # (..., D) = (..., D) / (..., 1)
    axis_vector = axis_vector * torch.reciprocal(axis_norm)

    # (...) or (..., 1)
    sin_phi = torch.sum(axis_vector*cross_ab, dim=-1, keepdim=keepdim)
    cos_phi = torch.sum(vec_a*vec_b, dim=-1, keepdim=keepdim)

    return torch.atan2(sin_phi, cos_phi)

def calc_torsion_nopbc(position_a: Tensor,
                       position_b: Tensor,
                       position_c: Tensor,
                       position_d: Tensor,
                       keepdim: bool = False,
                       ) -> Tensor:
    r"""Compute torsion angle :math:`A-B-C-D` formed by four positions :math:`A`, :math:`B`, :math:`C` and :math:`D`
        without periodic boundary condition.

    """

    vec_ba = calc_vector_nopbc(position_b, position_a)
    vec_cd = calc_vector_nopbc(position_c, position_d)
    vec_bc = calc_vector_nopbc(position_b, position_c)
    return calc_torsion_by_vectors(vec_ba, vec_cd, axis_vector=vec_bc, keepdim=keepdim)

def calc_torsion_pbc(position_a: Tensor,
                     position_b: Tensor,
                     position_c: Tensor,
                     position_d: Tensor,
                     pbc_box: Tensor,
                     keepdim: bool = False,
                     ) -> Tensor:
    r"""Compute torsion angle :math:`A-B-C-D` formed by four positions :math:`A`, :math:`B`, :math:`C` and :math:`D`
        at periodic boundary condition.

    """

    vec_ba = calc_vector_pbc(position_b, position_a, pbc_box)
    vec_cd = calc_vector_pbc(position_c, position_d, pbc_box)
    vec_bc = calc_vector_pbc(position_b, position_c, pbc_box)
    return calc_torsion_by_vectors(vec_ba, vec_cd, axis_vector=vec_bc, keepdim=keepdim)

def calc_torsion(position_a: Tensor,
                 position_b: Tensor,
                 position_c: Tensor,
                 position_d: Tensor,
                 pbc_box: Tensor = None,
                 keepdim: bool = False,
                 ) -> Tensor:

    r"""Compute torsion angle :math:`A-B-C-D` formed by four positions :math:`A`, :math:`B`, :math:`C` and :math:`D`
        with or without periodic boundary condition.

    """

    if pbc_box is None:
        return calc_torsion_nopbc(
            position_a, position_b, position_c, position_d, keepdim=keepdim)

    return calc_torsion_pbc(
        position_a, position_b, position_c, position_d, pbc_box=pbc_box, keepdim=keepdim)

def coulomb_interaction(q_i: Tensor,
                        q_j: Tensor,
                        r_ij: Tensor,
                        mask: Tensor = None,
                        coulomb_const: float = 1,
                        ):
    r"""Calculate Coulomb interaction.

    """

    energy = coulomb_const * q_i * q_j / r_ij

    if mask is None:
        return energy

    return energy * mask

def lennard_jones_potential(epsilon: Tensor, sigma: Tensor, r_ij: Tensor, mask: Tensor = None) -> Tensor:
    r"""Calculate Lennard-Jones (LJ) potential with :math:`\epsilon` and :math:`\sigma`.

    """
    # (\sigma / r_{ij}) ^ 6
    r0_6 = torch.pow(sigma / r_ij, 6)
    # 4 * \epsilon * (\sigma / r_{ij}) ^ 6
    ene_bcoeff = 4 * epsilon * r0_6
    # 4 * \epsilon * (\sigma / r_{ij}) ^ 12
    ene_acoeff = ene_bcoeff * r0_6

    energy = ene_acoeff - ene_bcoeff

    if mask is None:
        return energy
    return energy * mask

def lennard_jones_potential2(epsilon: Tensor, r_0: Tensor, r_ij: Tensor, mask: Tensor = None) -> Tensor:
    r"""Calculate Lennard-Jones (LJ) potential with :math:`\epsilon` and :math:`r_0`.

    """
    # (\r_0 / r_{ij}) ^ 6
    r0_6 = torch.pow(r_0 / r_ij, 6)
    # 2 * \epsilon * (r_0 / r_{ij}) ^ 6
    ene_bcoeff = 2 * epsilon * r0_6
    # \epsilon * (r_0 / r_{ij}) ^ 12
    ene_acoeff = epsilon * r0_6 * r0_6

    energy = ene_acoeff - ene_bcoeff

    if mask is None:
        return energy
    return energy * mask


def get_integer(value: Union[int, Tensor, Parameter, ndarray]) -> int:
    
    if value is None:
        return None
    if isinstance(value, Tensor):
        value = value.cpu().numpy()
    return int(value)


def get_ndarray(value: Union[Tensor, Parameter, ndarray, List[float], Tuple[float]],
                dtype: type = None) -> ndarray:
    r"""get ndarray type of the input value

    """
    if value is None:
        return None
    if isinstance(value, (Tensor, Parameter)):
        value = value.numpy()
        if dtype is not None:
            value = value.astype(dtype)
    else:
        value = np.array(value, dtype)
    return value


def get_tensor(value: Union[float, int, Tensor, Parameter, ndarray, List[float], Tuple[float]],
               dtype: type = None,
               device: str = None) -> Tensor:
    r"""get mindspore.Tensor type of the input value

    """
    if value is None:
        return None
    if isinstance(value, ndarray):
        value = torch.as_tensor(value, device=device)
    elif isinstance(value, Parameter):
        value = value.clone()
    elif not isinstance(value, Tensor):
        value = torch.tensor((value,), device=device)

    return value.to(dtype)

def get_ms_array(value: Union[float, int, Tensor, Parameter, ndarray, list, tuple],
                 dtype: type = None,
                 device: str = None,
                 ) -> Union[Tensor, Parameter]:
    r"""get mindspore.Tensor type of the input value

    """

    if value is None:
        return None
    if isinstance(value, ndarray):
        value = torch.as_tensor(value, dtype=dtype, device=device)
    elif not isinstance(value, (Tensor, Parameter)):
        (value,) = torch.tensor((value,), dtype=dtype, device=device)

    return value


def check_broadcast(shape0: tuple, shape1: tuple) -> tuple:
    r"""Check whether the two shapes match the rule of broadcast.

    """
    if shape0 is None:
        return shape1
    if shape1 is None:
        return shape0

    if len(shape0) < len(shape1):
        shape0 = (1,) * (len(shape1) - len(shape0)) + shape0
    if len(shape0) > len(shape1):
        shape1 = (1,) * (len(shape0) - len(shape1)) + shape1

    shape = ()
    for s0, s1 in zip(shape0, shape1):
        if s0 == s1:
            s = s0
        else:
            if s0 == 1:
                s = s1
            elif s1 == 1:
                s = s0
            else:
                raise ValueError(f'{shape0} and {shape1} cannot be broadcast to each other!')
        shape += (s,)
    return shape


def any_none(iterable: Iterable) -> bool:
    r"""Return True if ANY values x in the iterable is None.

    """
    return any([i is None for i in iterable])


def all_none(iterable: Iterable) -> bool:
    r"""Return True if ALL values `x` in the `iterable` is None..

    """
    return all([i is None for i in iterable])


def any_not_none(iterable: Iterable) -> bool:
    r"""Return True if ANY values `x` in the `iterable` is NOT None.

    """
    return any([i is not None for i in iterable])


def all_not_none(iterable: Iterable) -> bool:
    r"""Return True if ALL values `x` in the `iterable` is Not None.

    """
    return all([i is not None for i in iterable])


def get_arguments(locals_: dict, kwargs: dict = None) -> dict:
    r"""get arguments of a class

    Args:
        locals_ (dict): Dictionary of the arguments from `locals()`.
        kwargs (dict): Dictionary of keyword arguments (kwargs) of the class.

    Returns:
        args (dict): Dictionary of arguments

    Supported Platforms:
        ``Ascend`` ``GPU`` ``CPU``

    """

    if '__class__' in locals_.keys():
        locals_.pop('__class__')

    arguments = {}
    if 'self' in locals_.keys():
        cls = locals_.pop('self')
        arguments['cls_name'] = cls.__class__.__name__

    def _set_arguments(args_: dict):
        def _convert(value):
            if value is None or isinstance(value, (int, float, bool, str,
                                                   time, timedelta, date)):
                return value
            if isinstance(value, ndarray):
                return value.tolist()
            if isinstance(value, (Tensor, Parameter)):
                return value.detach().cpu().numpy().tolist()
            if isinstance(value, (list, tuple)):
                return [_convert(v) for v in value]
            if isinstance(value, dict):
                if 'cls_name' in value.keys():
                    return value
                dict_ = value.copy()
                for k, v in value.items():
                    dict_[k] = _convert(v)
                return dict_

            cls_name = value.__class__.__name__
            if hasattr(value, '_kwargs'):
                value = value.__dict__['_kwargs']
            elif hasattr(value, 'init_args'):
                value = value.__dict__['init_args']
            else:
                value = value.__class__.__name__

            if isinstance(value, dict) and 'cls_name' not in value.keys():
                dict_ = {'cls_name': cls_name}
                dict_.update(_set_arguments(value))
                value = dict_

            return value

        for k, v in args_.items():
            args_[k] = _convert(v)
        return args_

    kwargs_ = {}
    if 'kwargs' in locals_.keys():
        kwargs_: dict = locals_.pop('kwargs')

    if kwargs is None:
        kwargs = kwargs_

    if 'cls_name' in kwargs.keys():
        kwargs.pop('cls_name')

    arguments.update(_set_arguments(locals_))
    arguments.update(_set_arguments(kwargs))

    return arguments

initilizer = torch.nn.init.__dir__()[18:]
def get_initializer(cls_name: Union[str, Tensor], in_features, out_features=None, **kwargs):
    r"""get initializer by name

    Args:
        cls_name (Union[Initializer, str, dict, Tensor]): Class name of Initializer.
        kwargs (dict): Dictionary of keyword arguments (kwargs) of the class.

    Returns:
        Initializer

    Supported Platforms:
        ``Ascend`` ``GPU`` ``CPU``

    """
    if isinstance(cls_name, str):
        if cls_name[-1] != '_':
            cls_name += '_'
        if cls_name.lower() not in initilizer:
            raise ValueError(f"For 'initializer', the class corresponding to '{cls_name}' was not found.")
        else:
            init = torch.nn.init.__getattribute__(cls_name)
            if out_features == None:
                return init(torch.empty(in_features, ))
            else:
                return init(torch.empty(in_features, out_features))

    if isinstance(cls_name, (Tensor, Parameter, ndarray)):
        return get_tensor(cls_name, torch.float32)

    if isinstance(cls_name, dict):
        return get_initializer(**cls_name)

    raise TypeError(f'The cls_name must be str, dict or Tensor but got: {type(cls_name)}')


def _bonds_in(bonds, bond):
    """ Check bonds exists in both sets. """
    return (bonds == bond).all(-1)


def bonds_in(bonds, batch_bond):
    """ Return if batch_bond exists in bonds.
    Args:
        bonds: The total bonds set.
        batch_bond: The input bond set.

    Returns:
        If batch_bond exists in bonds, the mask will be 1, else 0.
    """
    return torch.vmap(_bonds_in, in_dims=(None, -2), out_dims=0)(bonds, batch_bond).sum(dim=-3)

class torchSpongeDevice:
    def __init__(self, device) -> None:
        self.device = device

    def set_device(self, device):
        self.device = device
    def __call__(self):
        return self.device
    
GLOBAL_DEVICE_SPONGE = torchSpongeDevice('cpu')

def set_torchsponge_global_device(device):
    global GLOBAL_DEVICE_SPONGE
    GLOBAL_DEVICE_SPONGE.set_device(device)

def toxyz(atom, data, box):
    num = len(data)
    box = box.reshape(-1)
    content = f'{num}\nLattice="{box[0]} 0.0 0.0 0.0 {box[1]} 0.0 0.0 0.0 {box[2]}" Properties=species:S:1:pos:R:3 pbc=T T T\n'
    atom_block = '\n'.join(f'{a}     {x:.8f}     {y:.8f}     {z:.8f}' for a, (x, y, z) in zip(atom, data))
    return content+atom_block